import tensorflow as tf
from tensorflow.examples.tutorials.mnist import input_data

mnist = input_data.read_data_sets("input_data/", one_hot=True)
# 超参数
learning_rate = 0.01
epochs = 1000
display_step = 20
batch_size = 100

# 定义占位符
X = tf.placeholder(tf.float32, [None, 784])
Y = tf.placeholder(tf.float32, [None, 10])

# 定义变量
W = tf.Variable(tf.zeros([784, 10]))
b = tf.Variable(tf.zeros([10]))

# 定义模型
Wx_plus_b = tf.matmul(X, W) + b
pred = tf.nn.softmax(Wx_plus_b)
# 定义损失函数与梯度下降公式
loss = tf.reduce_mean(-tf.reduce_sum(Y*tf.log(pred), reduction_indices=1))
optimizer = tf.train.GradientDescentOptimizer(learning_rate=learning_rate).minimize(loss)
init = tf.global_variables_initializer()

# 启动tf系统流程
with tf.Session() as sess:
    sess.run(init)
    for epoch in range(epochs):
        avg_loss = 0.
        total_batch = int(mnist.train.num_examples / batch_size)
        for i in range(total_batch):
            batch_x, batch_y = mnist.train.next_batch(batch_size)
            _, c = sess.run([optimizer, loss], feed_dict={X:batch_x, Y:batch_y})
            avg_loss += c/total_batch

        # 输出当前的损失函数
        # if (epoch+1) % display_step == 0:
        #     print("Epoch:", '%04d' % (epoch+1), "cost=", "{:.9f}".format(avg_loss))
        print("Epoch:", '%04d' % (epoch+1), "cost=", "{:.9f}".format(avg_loss))
    print("Optimizer Finished")

    # 验证模型的准确率
    correct_pred = tf.equal(tf.argmax(pred, 1), tf.argmax(Y,1))
    accuracy = tf.reduce_mean(tf.cast(correct_pred, tf.float32))
    print("Accuracy:",accuracy.eval({X:mnist.test.images, Y:mnist.test.labels}))
